#include "traj_server/traj_server.h"
#include <rosbag/bag.h>
#include <rosbag/view.h>
#include "transform_utils/type_transform.h"
#include "transform_utils/display.h"

namespace Planner
{
	using std::string;
	using std::vector;

	TrajServer::TrajServer(const ros::NodeHandle &node_handle)
		: nh_(node_handle)
	{
		init();
	}

	void TrajServer::init()
	{
		string traj_bag, traj_topic;
		nh_.param("traj_server/traj_bag", traj_bag, string(""));
		nh_.param("traj_server/traj_topic", traj_topic, string(""));
		nh_.param("traj_server/hover_time", hover_time_, 5.0);
		nh_.param("traj_server/replan_time_window", replan_time_window_, 5.0);
		nh_.param("traj_server/replan_time_cost", replan_time_cost_, 0.5);
		nh_.param("traj_server/visualization", visualization_, false);
		nh_.param("traj_server/replan_triggered", replan_triggered_, false);

		vector<quad_msgs::Trajectory::Ptr> traj_data;
		try
		{
			rosbag::Bag bag;
			bag.open(traj_bag, rosbag::bagmode::Read);
			rosbag::View view(bag, rosbag::TopicQuery(traj_topic));

			for (rosbag::MessageInstance const m : view)
			{
				quad_msgs::Trajectory::Ptr s = m.instantiate<quad_msgs::Trajectory>();
				if (s != nullptr)
					traj_data.push_back(s);
			}
		}
		catch (const std::exception &e)
		{
			std::cerr << e.what() << '\n';
			status_.mode = INIT_FAIL;
		}

		if (traj_data.size() != 1)
		{
			ROS_ERROR("Global trajectory msgs should obtain 1 msg, but now msg size is %d", traj_data.size());
			status_.mode = INIT_FAIL;
		}
		else
		{
			global_traj_msg_ = traj_data[0];
			status_.mode = INIT_SUCCESS;
			status_.curr_seg = 0;
			// printTrajectoryMsg(*global_traj_msg_);
		}

		nh_.param("traj_server/control_rate", control_rate_, 10.0);
		nh_.param("traj_server/replan_rate", replan_rate_, 1.0);

		trigger_server_ = nh_.advertiseService("trigger_traj_server", &TrajServer::triggerCallback, this);
		control_command_pub_ = nh_.advertise<quad_msgs::ControlCommand>("/control_command", 5);
		replan_client_ = nh_.serviceClient<quad_msgs::ReplanCheck>("replan_check");

		global_traj_pub_ = nh_.advertise<nav_msgs::Path>("global_traj_vis", 1);
		replan_request_traj_pub_ = nh_.advertise<nav_msgs::Path>("replan_request_traj_vis", 1);
		replan_response_traj_pub_ = nh_.advertise<nav_msgs::Path>("replan_response_traj_vis", 1);

		if (visualization_)
		{
			displayTrajectory(*global_traj_msg_, global_traj_pub_);
			ros::Duration(1.0).sleep();
			displayTrajectory(*global_traj_msg_, global_traj_pub_);
			ROS_INFO("Global trajectory visualized");
		}
	}

	// 使用智能指针实现global_traj_msg_的读写锁
	void TrajServer::controlCommandCallback(const ros::TimerEvent &event)
	{
		// 把可能存在竞态的数据一次性读取出来
		quad_msgs::Trajectory::Ptr trajectory_data_ptr;
		Status status;
		{
			std::lock_guard<std::mutex> guard(mutex_);
			trajectory_data_ptr = global_traj_msg_;
			status = status_;
		}

		switch (status.mode)
		{
		case INIT_FAIL:
		{
			ROS_ERROR("Initialized failed");
			break;
		}

		case INIT_SUCCESS:
		{
			// 初始化过程，使无人机在轨迹起始位置悬停
			if ((ros::Time::now() - status.init_time).toSec() <= hover_time_)
			{
				quad_msgs::State s = trajectory_data_ptr->waypoints.front();
				status.last_command.position = s.position;
				status.last_command.velocity = s.velocity;
				status.last_command.acceleration = s.acceleration;
				status.last_command.command_flag = quad_msgs::ControlCommand::HOVER;
				control_command_pub_.publish(status.last_command);
			}
			else
			{
				status.init_time = ros::Time::now();
				status.mode = SAFE_FLY;
				status.last_time_seq = 0;
				control_command_pub_.publish(status.last_command);
			}
			break;
		}

		case SAFE_FLY:
		{
			double t = (ros::Time::now() - status.init_time).toSec();
			// 判断当前时刻对应分段轨迹的段数
			while (t - status.last_time_seq > trajectory_data_ptr->polynomials[status.curr_seg].duration)
			{
				status.last_time_seq += trajectory_data_ptr->polynomials[status.curr_seg].duration;
				status.curr_seg++;
				// 抵达终点，更改state
				if (status.curr_seg == trajectory_data_ptr->segments)
				{
					status.mode = FINISH_TASK;
					status.curr_seg--;
					status.last_time_seq -= trajectory_data_ptr->polynomials[status.curr_seg].duration;
					ROS_INFO("Finish the whole trajectory");
					break;
				}
			}
			double duration = t - status.last_time_seq;
			if (status.mode == FINISH_TASK)
				duration = trajectory_data_ptr->polynomials.back().duration;

			// 计算所需目标点
			vector<double> coeffsx = trajectory_data_ptr->polynomials[status.curr_seg].coeffs_x;
			vector<double> coeffsy = trajectory_data_ptr->polynomials[status.curr_seg].coeffs_y;
			vector<double> coeffsz = trajectory_data_ptr->polynomials[status.curr_seg].coeffs_z;
			quad_msgs::State s = getStateFromCoeffs(coeffsx, coeffsy, coeffsz, duration);
			status.last_command.position = s.position;
			status.last_command.velocity = s.velocity;
			status.last_command.acceleration = s.acceleration;
			status.last_command.command_flag = quad_msgs::ControlCommand::FLY;
			control_command_pub_.publish(status.last_command);
			break;
		}

		case FINISH_TASK:
		{
			status.last_command.command_flag = quad_msgs::ControlCommand::HOVER;
			control_command_pub_.publish(status.last_command);
			break;
		}

		case REPLAN_FAIL:
		{
			status.last_command.command_flag = quad_msgs::ControlCommand::BRAKE;
			control_command_pub_.publish(status.last_command);
			break;
		}

		default:
			break;
		}

		// 更新status
		{
			std::lock_guard<std::mutex> guard(mutex_);
			status_ = status;
		}
	}

	void TrajServer::replanCheckCallback(const ros::TimerEvent &event)
	{
		// 先读信息
		quad_msgs::Trajectory::Ptr trajectory_data_ptr;
		Status status;
		{
			std::lock_guard<std::mutex> guard(mutex_);
			trajectory_data_ptr = global_traj_msg_;
			status = status_;
		}

		// 只有无人机正常飞行的时候才需要进行局部规划
		// TODO 若上次规划失败应该加上重新进行局部规划尝试的功能
		if (status.mode != SAFE_FLY)
			return;

		// 构造replan_check_srv的request部分
		quad_msgs::ReplanCheck replan_check_srv;
		double t = (ros::Time::now() - status.init_time).toSec();
		double curr_time_stamp = status.last_time_seq;
		// NOTE 这里明确一下，i指代的是第i个路径点是否可以被加入到request中
		for (int i = status.curr_seg + 1; i <= trajectory_data_ptr->segments; i++)
		{
			curr_time_stamp += trajectory_data_ptr->polynomials[i - 1].duration; // 表示第i个路径点对应的时间戳
			// 选择进行局部规划的路径点距离当前时刻的时间间隔不能少于一次局部规划的最大耗时
			if (curr_time_stamp - t < replan_time_cost_)
				continue;
			// 只对一定时间滑动窗口内的路径点进行局部规划
			if (curr_time_stamp - t < replan_time_window_)
			{
				if (!replan_check_srv.request.waypoints.empty())
					replan_check_srv.request.polynomials.push_back(trajectory_data_ptr->polynomials[i - 1]);
				replan_check_srv.request.waypoints.push_back(trajectory_data_ptr->waypoints[i]);
				replan_check_srv.request.waypoint_id.push_back(i);
			}
			else
				break;
		}

		// 离终点很近，剩余路径点均在一个小时间段内，不进行局部规划，认为之前的局部规划保障了剩余轨迹的安全性
		if (replan_check_srv.request.polynomials.empty())
			return;

		// 判断待重规划路径是否包含终点，这与局部规划区间的终点是否允许移动相关
		if (replan_check_srv.request.waypoint_id.back() == trajectory_data_ptr->segments)
			replan_check_srv.request.obtain_goal = true;
		else
			replan_check_srv.request.obtain_goal = false;

		if (visualization_)
			displayTrajectory(replan_check_srv.request, replan_request_traj_pub_);

		// 请求轨迹局部规划服务
		if (!replan_client_.call(replan_check_srv))
		{
			ROS_ERROR("Replan server is not working");
			return;
		}

		if (visualization_)
			displayTrajectory(replan_check_srv.response, replan_response_traj_pub_);

		// 首先释放用于只读的trajectory_data_ptr
		trajectory_data_ptr.reset(new quad_msgs::Trajectory);

		switch (replan_check_srv.response.replan_flag)
		{
		case quad_msgs::ReplanCheck::Response::SAFE:
		{
			ROS_INFO("Trajectory in sliding window is safe");
			break;
		}

		case quad_msgs::ReplanCheck::Response::REPLAN_SUCCESS:
		{
			ROS_INFO("Replan triggered, trajectory is revised safely");
			// 写入轨迹信息
			std::lock_guard<std::mutex> guard(mutex_);

			// 如果还有其它ptr指向global_traj_msg，即其它线程正在进行读操作，则拷贝副本后修改
			if (!global_traj_msg_.unique())
				global_traj_msg_.reset(new quad_msgs::Trajectory(*global_traj_msg_));

			for (int i = replan_check_srv.response.start_id + 1; i < replan_check_srv.response.end_id; i++)
				global_traj_msg_->waypoints[i] = replan_check_srv.response.waypoints[i - replan_check_srv.response.start_id - 1];
			
			for (int i = replan_check_srv.response.start_id; i < replan_check_srv.response.end_id; i++)
				global_traj_msg_->polynomials[i] = replan_check_srv.response.polynomials[i - replan_check_srv.response.start_id];

			break;
		}

		case quad_msgs::ReplanCheck::Response::MOVE_END:
		{
			ROS_INFO("Replan triggered, end of the sliding window is revised");
			// 先更新滑动窗口内轨迹信息
			std::lock_guard<std::mutex> guard(mutex_);

			// 如果还有其它ptr指向global_traj_msg，即其它线程正在进行读操作，则拷贝副本后修改
			if (!global_traj_msg_.unique())
				global_traj_msg_.reset(new quad_msgs::Trajectory(*global_traj_msg_));

			for (int i = replan_check_srv.response.start_id + 1; i <= replan_check_srv.response.end_id; i++)
				global_traj_msg_->waypoints[i] = replan_check_srv.response.waypoints[i - replan_check_srv.response.start_id - 1];
			
			for (int i = replan_check_srv.response.start_id; i < replan_check_srv.response.end_id; i++)
				global_traj_msg_->polynomials[i] = replan_check_srv.response.polynomials[i - replan_check_srv.response.start_id];
			break;
		}

		case quad_msgs::ReplanCheck::Response::REPLAN_FAIL:
		{
			ROS_ERROR("Replan failed, call emergency brake");
			std::lock_guard<std::mutex> guard(mutex_);
			status_.mode = REPLAN_FAIL;
		}

		default:
			break;
		}
	}

	bool TrajServer::triggerCallback(quad_msgs::Trigger::Request &request,
									 quad_msgs::Trigger::Response &response)
	{
		if (request.trigger == true)
		{
			control_command_timer_ = nh_.createTimer(ros::Duration(1 / control_rate_), &TrajServer::controlCommandCallback, this);
			if (replan_triggered_)
				replan_check_timer_ = nh_.createTimer(ros::Duration(1 / replan_rate_), &TrajServer::replanCheckCallback, this);
			status_.init_time = ros::Time::now();
			response.ack = true;
		}

		return true;
	}
} // namespace Planner